# -*- coding: utf-8 -*-
"""
Spyder Editor

This is a temporary script file.
"""

import torch
from torch.autograd import Variable

x_data = Variable(torch.Tensor([[1.], [2.], [3.], [4.]]))
y_data = Variable(torch.Tensor([[0.], [0.], [1.], [1.]]))

class Model(torch.nn.Module):
    def __init__(self):
        '''
        In the constructor we instantiate two nn.Linear module
        '''
        super(Model, self).__init__()
        self.linear = torch.nn.Linear(1, 1) # One in and one out
        self.sigmoid = torch.nn.Sigmoid()
        
    def forward(self, x):
        '''
        In the forward function we accept a Variable of input data and we
        must return a Variable of output data. We can use Modules defined 
        in the constructor as well as arbitrary operators on Variables.
        '''
        y_pred = self.sigmoid(self.linear(x))
        return y_pred
    
# our model
model = Model()

'''
Constructor our loss function and an Optimizer. The call to model.parameters()
in the SGD constructor will contain the learnable parameters of the two 
nn.Linear modules which are members of the model.
'''
criterion = torch.nn.BCELoss(size_average=True)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# Training loop
for epoch in range(500):
    # Forward pass: compute predicted y by passing x to the model
    y_pred = model(x_data)
    
    # Compute and print loss
    loss = criterion(y_pred, y_data)
    print(epoch, loss.data[0])
    
    # Zero gradients, perform a backward pass, and update the weights.
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
# After training
hour_var = Variable(torch.Tensor([[0.5]]))
print("predict (after training)", 0.5, model.forward(hour_var).data[0][0])
hour_var = Variable(torch.Tensor([[7.]]))
print("predict (after training)", 7., model.forward(hour_var).data[0][0])
